// Copyright (C) joyinter
// All rights reserved.
// 
// Author : Jiang Li
// Date   : 2005-5-8

#include "stdafx.h"
#include "checkpe.h"
#include "crc.h"
#include "apihelp.h"

namespace JoyGuard
{

//##ModelId=44C4B60E017A
CheckPe::CheckPe()
{
    // process data
    m_process   = NULL;
    m_processId = 0;
    m_module    = NULL;
    m_checkRva  = 0;
    m_checkSize = 0;

    // pe file data
    m_fileHash  = 0;

    // check control
    m_error     = 0;
    m_waitEvent = NULL;
    m_readSize  = 64*1024;
    m_readDelay = 100;
    m_readed    = 0;
    m_stopped   = false;
}

//##ModelId=44C4B60E017B
bool CheckPe::ReadCodeData(const char* fileName,
    std::vector<BYTE>& codeData, DWORD& codeRva, DWORD& codeSize, DWORD& imageSize)
{
    // read data from pe file
    HANDLE peFile = ::CreateFile(fileName, GENERIC_READ,
        FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if(peFile == INVALID_HANDLE_VALUE)
    {
        return false;
    }

    DWORD dwSize = ::GetFileSize(peFile, NULL);
    DWORD dwReaded = 0;
    std::vector<BYTE> peBuf(dwSize);
    if(!::ReadFile(peFile, &peBuf[0], dwSize, &dwReaded, NULL))
    {
        return false;
    }

    // analyze image.

    BYTE* baseAddr = &peBuf[0];

    // get dos header
    IMAGE_DOS_HEADER* dosHeader    = (IMAGE_DOS_HEADER*)baseAddr;
    if(dosHeader->e_magic != IMAGE_DOS_SIGNATURE)
    {
        return NULL;
    }
    if(dosHeader->e_lfanew <= 0)
    {
        return false;
    }

    // sectName is ".text"
    char sectName[8];
    sectName[1] = 't';
    sectName[4] = 's';

    // get nt headers
    IMAGE_NT_HEADERS* ntheaders   = (IMAGE_NT_HEADERS*)(baseAddr+dosHeader->e_lfanew);
    if(ntheaders->Signature != IMAGE_NT_SIGNATURE)
    {
        return false;
    }

    sectName[0] = '.';
    sectName[2] = 'f';

    // find a section match the sectName
    int sectCount = ntheaders->FileHeader.NumberOfSections;
    IMAGE_SECTION_HEADER* sectheaders = (IMAGE_SECTION_HEADER*)(ntheaders+1);
    IMAGE_SECTION_HEADER* foundHeader = NULL;
    sectName[3] = 'x';
    sectName[5] = 0;
    ++ sectName[4];
    -- sectName[2];
    for(int i = 0; i < (int)sectCount; i ++)
    {
        IMAGE_SECTION_HEADER* header = sectheaders+i;
        if(header->Characteristics == 0)
        {
            return false;
        }
        if(strcmp((char*)header->Name, sectName) == 0)
        {
            foundHeader = header;
            break;
        }
    }
    if(foundHeader == NULL)
    {
        return false;
    }

    const BYTE* sectBuf = baseAddr+foundHeader->PointerToRawData;
    codeData.assign(sectBuf, sectBuf+foundHeader->Misc.VirtualSize);
    codeRva     = foundHeader->VirtualAddress;
    codeSize    = foundHeader->Misc.VirtualSize;
    imageSize   = ntheaders->OptionalHeader.SizeOfImage;

    return true;
}

//##ModelId=44C4B60E018A
bool CheckPe::CreateByFile(const char* fileName, DWORD hash)
{
    if(!ReadCodeData(fileName, m_baseData, m_checkRva, m_checkSize, m_imageSize))
    {
        return false;
    }
    if(hash != 0)
    {
        Crc csum;
        csum.Update(&m_baseData[0], (int)m_baseData.size());
        if(hash != csum.Final())
        {
            m_error = EC_CHECK_FILE;
            return false;
        }
    }

    const char* modName = fileName + ::strlen(fileName) -1;
    for(; modName > fileName; -- modName)
    {
        if(*modName == '\\' || *modName == '/')
        {
            ++ modName;
            break;
        }
    }

    return CreateByModule(modName);
}

#define IMAGE_REL_BASED_SECTION 6
#define IMAGE_REL_BASED_REL32 7
#define IMAGE_REL_BASED_HIGH3ADJ 11

//##ModelId=44C4B60E0197
bool CheckPe::CreateByModule(const char* modName)
{
    if(m_process == NULL || m_processId == 0)
    {
        return false;
    }

    HANDLE hModuleSnap = INVALID_HANDLE_VALUE;
    MODULEENTRY32 me32;
    Sleep(1000);
    hModuleSnap = ApiHelp::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, m_processId);

    // Set the size of the structure before using it.
    me32.dwSize = sizeof( MODULEENTRY32 );

    if(!ApiHelp::Module32First(hModuleSnap, &me32))
    {
        CloseHandle(hModuleSnap);
        return false;
    }

    do
    {
        if(stricmp(me32.szModule, modName) == 0)
        {
            m_module = me32.modBaseAddr;
            break;
        }
    }
    while(ApiHelp::Module32Next(hModuleSnap, &me32));

    ::CloseHandle(hModuleSnap);

    if(m_module == NULL)
    {
        return false;
    }

    return Create();
}

//##ModelId=44C4B60E0199
bool CheckPe::Create()
{
    CheckStop();

    std::vector<BYTE> imageTemp(m_imageSize);
    DWORD readDelay = m_readDelay;
    m_readDelay = 0;
    ReadMemory(&imageTemp[0], m_module, m_imageSize);
    m_readDelay = readDelay;

    BYTE* baseAddr = &imageTemp[0];
    IMAGE_DOS_HEADER* dosHeader = (IMAGE_DOS_HEADER*)baseAddr;
    IMAGE_NT_HEADERS* ntHeaders = (IMAGE_NT_HEADERS*)(baseAddr+dosHeader->e_lfanew);

    // do relocation
    BYTE* baseData = &m_baseData[0];
    DWORD Diff = (DWORD)((UINT_PTR)m_module - ntHeaders->OptionalHeader.ImageBase);
    IMAGE_BASE_RELOCATION* baseRel = (IMAGE_BASE_RELOCATION*)
        (baseAddr + ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
    for(; baseRel->VirtualAddress;
        baseRel = (IMAGE_BASE_RELOCATION*)((BYTE*)baseRel + baseRel->SizeOfBlock))
    {
        if( baseRel->VirtualAddress < m_checkRva ||
            baseRel->VirtualAddress > m_checkRva + m_checkSize )
        {
            continue;
        }
        LONG Temp;
        LONGLONG Temp64;
        DWORD VA = baseRel->VirtualAddress - m_checkRva;
        DWORD count = (baseRel->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION))/sizeof(WORD);
        WORD* pType = (WORD*)(baseRel+1);
        WORD* pTypeMax = pType + count;

        for(; pType < pTypeMax; ++ pType)
        {
            PVOID FixupVA = baseData + VA + (*pType & 0xFFF);
            WORD fixType = *pType >> 12;
            switch(fixType)
            {
            case IMAGE_REL_BASED_HIGHLOW :
                *(LONG*)FixupVA += (ULONG) Diff;
                break;

            case IMAGE_REL_BASED_HIGH :
                Temp = *(PUSHORT)FixupVA << 16;
                Temp += (ULONG) Diff;
                *(PUSHORT)FixupVA = (USHORT)(Temp >> 16);
                break;

            case IMAGE_REL_BASED_HIGHADJ :
                Temp = *(PUSHORT)FixupVA << 16;
                ++pType;
                Temp += (LONG)(*(PSHORT)pType);
                Temp += (ULONG) Diff;
                Temp += 0x8000;
                *(PUSHORT)FixupVA = (USHORT)(Temp >> 16);
                break;

            case IMAGE_REL_BASED_LOW :
                Temp = *(PSHORT)FixupVA;
                Temp += (ULONG) Diff;
                *(PUSHORT)FixupVA = (USHORT)Temp;
                break;

            case IMAGE_REL_BASED_IA64_IMM64:

                break;

            case IMAGE_REL_BASED_DIR64:

                *(ULONG_PTR*)FixupVA += Diff;

                break;

            case IMAGE_REL_BASED_MIPS_JMPADDR :
                Temp = (*(PULONG)FixupVA & 0x3ffffff) << 2;
                Temp += (ULONG) Diff;
                *(PULONG)FixupVA = (*(PULONG)FixupVA & ~0x3ffffff) |
                                                ((Temp >> 2) & 0x3ffffff);

                break;

            case IMAGE_REL_BASED_ABSOLUTE :
                //
                // Absolute - no fixup required.
                //
                break;

            case IMAGE_REL_BASED_SECTION :
                //
                // Section Relative reloc.  Ignore for now.
                //
                break;

            case IMAGE_REL_BASED_REL32 :
                //
                // Relative intrasection. Ignore for now.
                //
                break;

           case IMAGE_REL_BASED_HIGH3ADJ :
               Temp64 = *(PUSHORT)FixupVA << 16;
               ++pType;
               Temp64 += (LONG)((SHORT)pType[1]);
               Temp64 <<= 16;
               Temp64 += (LONG)((USHORT)pType[0]);
               Temp64 += Diff;
               Temp64 += 0x8000;
               Temp64 >>=16;
               Temp64 += 0x8000;
               *(PUSHORT)FixupVA = (USHORT)(Temp64 >> 16);
               ++pType;
               break;
            default :
                return false;
            } // switch
        }
    }

#if 0
    BYTE* src = &m_baseData[0];
    BYTE* dst = baseAddr+m_checkRva;
    for(int i = 0; i < (int)m_checkSize; ++ i)
    {
        if(src[i] != dst[i])
        {
            GetTickCount();
        }
    }

    int cc = memcmp(src, dst, m_checkSize);
#endif

    return true;
}

//##ModelId=44C4B60E019A
bool CheckPe::Check()
{
    return true;
}

//##ModelId=44C4B60E0169
void CheckPe::CheckStop()
{
    if(m_stopped)
    {
        throw ECheckStop();
    }
}

//##ModelId=44C4B60E016A
void CheckPe::ReadMemory(BYTE* dst, BYTE* src, int len)
{
    while(len > 0)
    {
        int c;
        if(len + m_readed >= m_readSize)
        {
            c = m_readSize-m_readed;
        }
        else
        {
            c = len;
        }
        SIZE_T readed = 0;
        if(!ApiHelp::ReadProcessMemory(m_process, src, dst, c, &readed))
        {
            m_error = EC_READ_MEMORY;
            throw ECheckStop();
        }
        m_readed += c;
        dst += c;
        src += c;
        len -= c;
        if(m_readed >= m_readSize)
        {
            m_readed -= m_readSize;
            ::WaitForSingleObject(m_waitEvent, m_readDelay);
            CheckStop();
        }
    }
}


} // namespace JoyGuard
